from openpose.src import networks
from estimator import TfPoseEstimator
from networks import get_graph_path

def human_score(human, num=18):
    '''
    计算一个human的整体score，如num-18时， 计算0~17的分数平均值，如果没有算0（因此会拉低分数）
    '''
    parts = human.body_parts
    score = 0
    for i in range(num):
        if i in parts:
            score += parts[i].score
    return float(score) / num

name2num = dict(
    Nose =0,
    Neck = 1,
    RShoulder = 2,
    RElbow = 3,
    RWrist = 4,
    LShoulder = 5,
    LElbow = 6,
    LWrist = 7,
    RHip = 8,
    RKnee = 9,
    RAnkle = 10,
    LHip = 11,
    LKnee = 12,
    LAnkle = 13,
    REye = 14,
    LEye = 15,
    REar = 16,
    LEar = 17,
    Background = 18,
)

num2name = {name2num[key]: key for key in name2num}


class OpenPoseEsitimator(object):
    def __init__(self, scales=None, w=256, h=256, model_name='mobilenet_thin'):
        self.scales = scales
        self.w = w
        self.h = h
        self.model_name = 'mobilenet_thin'
        self.e = TfPoseEstimator(get_graph_path(self.model_name), target_size=(self.w, self.h))
        

    def inference(self, image):
        '''
        只会推导出得分最大的那个，如果没有则返回None
        image是opencv的格式
        格式：名称->（宽上的坐标，高上的坐标）
        '''
        image_h = image.shape[0]
        image_w = image.shape[1]
        humans = self.e.inference(image, scales=self.scales)
        if len(humans) == 0:
            return None
        scores = [human_score(human) for human in humans]

        max_score = scores[0]
        max_human = humans[0]
        for i, score in enumerate(scores):
            if score > max_score:
                max_score = score
                max_human = humans[i]  
        ret = {num2name[num]: (int(max_human.body_parts[num].x * image_w + 0.5), int(max_human.body_parts[num].y * image_h + 0.5))
               for num in max_human.body_parts
              }
        return ret

    
    def inference_and_draw(self, image):
        '''
        只会推导出得分最大的那个，如果没有则返回None
        '''
        humans = self.e.inference(image, scales=self.scales)
        if len(humans) == 0:
            return None
        scores = [human_score(human) for human in humans]

        max_score = scores[0]
        max_human = humans[0]
        for i, score in enumerate(scores):
            if score > max_score:
                max_score = score
                max_human = humans[i]  
        return TfPoseEstimator.draw_humans(image, [max_human], imgcopy=False)
    